package com.gmrz.webauthn.protocol.v1.processor.attestation;

import co.nstant.in.cbor.model.DataItem;
import com.gmrz.uaf.common.Constants;
import com.gmrz.uaf.common.GuiceUtil;
import com.gmrz.uaf.crypto.CryptoEngine;
import com.gmrz.util.Convert;
import com.gmrz.webauthn.common.CBORDecodeUtil;
import com.gmrz.webauthn.common.LocationConstants;
import com.gmrz.webauthn.common.WebAuthnConstants;
import com.gmrz.webauthn.protocol.v1.json.WebAuthnSchemaBuilder;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnErrorCode;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnException;
import com.gmrz.webauthn.protocol.v1.schema.*;

import java.io.IOException;
import java.nio.charset.Charset;
import java.security.NoSuchAlgorithmException;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.List;
import java.util.Map;

import com.gmrz.webauthn.request.WebAuthnCredentialContext;
import org.apache.commons.codec.binary.Base64;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.bouncycastle.asn1.x500.RDN;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x500.style.BCStyle;
import org.bouncycastle.asn1.x500.style.IETFUtils;
import org.bouncycastle.cert.jcajce.JcaX509CertificateHolder;

public class PackedAttestationProcessor
        extends WebAuthnBaseAssertionProcessor
        implements IAttestationProcessor
{
    private static final Logger LOG = LogManager.getLogger(PackedAttestationProcessor.class);
    public static final String CERT_SUBJECT_FIELD_OU_LITERAL = "Authenticator Attestation";

    public PackedAttestationProcessor(CryptoEngine cryptoEngine){
        this.cryptoEngine = cryptoEngine;

    }

    public WebAuthnConstants.AttestationResult verify(WebAuthnCredentialContext publicKeyCredential, WebAuthenticator authenticator,boolean strictMode)
            throws WebAuthnException
    {
        if ((publicKeyCredential == null) || (null == publicKeyCredential.getResponse()))
        {
            String errorMsg = "public key credential is not present for attestation verification";
            LOG.error(errorMsg);
            return WebAuthnConstants.AttestationResult.UNAVAILABLE;
        }
        if (authenticator == null)
        {
            String errorMsg = "Authenticator is not present for attestation verification";
            LOG.error(errorMsg);
            return WebAuthnConstants.AttestationResult.UNAVAILABLE;
        }
        try
        {
            List<DataItem> dataitems = CBORDecodeUtil.cborDecode(Base64.decodeBase64(publicKeyCredential.getResponse().getAttestationObject()));
            co.nstant.in.cbor.model.Map cborMap = (co.nstant.in.cbor.model.Map) dataitems.get(0);

            java.util.Map<String, Object> map = CBORDecodeUtil.cborMapConversion(cborMap);

            byte[] rawAuthData = (byte[]) map.get("authData");

            AuthenticatorData authenticatorData = GuiceUtil.getProcessorInjector().getInstance(AuthenticatorData.class);
            authenticatorData.decode(rawAuthData);


            String formatName = (String) map.get("fmt");
            AttestationFormat format = AttestationFormat.forName(formatName);

            Object attStmtObj = map.get("attStmt");

            CredentialAttestation credentialAttestation = new CredentialAttestation().withAuthenticatorData(authenticatorData).withAuthenticatorDataBytes(rawAuthData).withAttestationStatement(attStmtObj).withFormat(format);
            credentialAttestation.setAuthenticatorDataBytes(rawAuthData);

            Map<String, Object> attStmt = (Map<String, Object>) credentialAttestation.getAttestationStatement();

            String clientDataStr = publicKeyCredential.getResponse().getClientDataJSON();
            String clientDataDecodedStr = new String(Base64.decodeBase64(clientDataStr), Charset.forName(Constants.UTF_8));
            CollectedClientData clientData = WebAuthnSchemaBuilder.getGson().fromJson(clientDataDecodedStr, CollectedClientData.class);

            if (attStmt.get("x5c") != null)
            {
                LOG.info("Starting basic attestation verification for packed format");
                List<Certificate> certs = fetchCertsFromAttestationStmt(attStmt);

                Certificate attestationCert = getAttestationCertificate(certs);
                verifyBasicAttestation(publicKeyCredential, attStmt, attestationCert,rawAuthData,clientDataDecodedStr.getBytes(),credentialAttestation);
                if (strictMode) {
                    validateCertificate(certs, WebAuthnConstants.getWebAuthenticatorSpec(authenticator.getAaguid().toString()), true);
                }
            }
            else if (attStmt.get("ecdaaKeyId") != null)
            {
                String errorMsg = "Unsupported attestation type";
                logAndThrow(errorMsg, WebAuthnErrorCode.STATUS_ATTESTATION_TYPE_NOT_SUPPORTED);
            }
            else
            {
                LOG.info("Starting self-attestaion verification for packed format");
                verifySelfAttestation(attStmt,rawAuthData,clientDataDecodedStr.getBytes(),credentialAttestation);
            }
        }
        catch (Exception e)
        {
            LOG.error("Error in attestation verification", e);
            return WebAuthnConstants.AttestationResult.FAILED;
        }
        LOG.info("Packed attestation verification status : " + this.attResult.name());
        return this.attResult;
    }

    private void verifySelfAttestation(Map<String, Object> attStmt,byte[] rawAuthData,byte[] clientData,CredentialAttestation credentialAttestation)
            throws WebAuthnException
    {
        try
        {
            int signAlgo = getSignAlgFromAttestaionMap(attStmt);
            JsonWebKey credentialPublicKey = credentialAttestation.getAuthenticatorData().getAttestation().getCredentialPublicKey();
            int credentialPublickKeyAlg = credentialPublicKey.getAlg().getSignAlgID();

            LOG.debug("Credential Public Key alg is [" + credentialPublickKeyAlg + "] & attestation statement alg is [" + signAlgo + "]");
            if (signAlgo != credentialPublickKeyAlg)
            {
                String errorMsg = "Credential Public Key alg does not match attestation statement alg";
                logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_SIGNATURE_VALIDATION_FAILED);
            }
            byte[] signedData = getSignedData(rawAuthData,clientData);

            byte[] attestationSignature = (byte[])attStmt.get("sig");
            if (attestationSignature == null)
            {
                String errorMsg = "Attestation signature not found";
                logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_SIGNATURE_VALIDATION_FAILED);
            }
            byte[] publicKey = Convert.fromBase64(credentialPublicKey.getPublicKeyInRawFormat());

            LOG.debug("Validating signature using alg:" + signAlgo);
            validateSignature(credentialPublicKey.getPublicKeyFormat(), publicKey, signAlgo, signedData, attestationSignature);
        }
        catch (Exception e)
        {
            String errorMsg = "Error in fetching signed data from attestation statement";
            logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_SIGNATURE_VALIDATION_FAILED, e);
        }
    }

    private void verifyBasicAttestation(WebAuthnCredentialContext authenticatorResponse, Map<String, Object> attStmt, Certificate attestationCert,byte[] rawAuthData,byte[] clientData,CredentialAttestation credentialAttestation)
            throws IOException, NoSuchAlgorithmException, WebAuthnException
    {
        byte[] signedData = getSignedData(rawAuthData,clientData);

        byte[] attestationSignature = (byte[])attStmt.get("sig");

        int signAlgo = getSignAlgFromAttestaionMap(attStmt);

        LOG.debug("Validating signature using alg:" + signAlgo);
        validateSignature(signAlgo, attestationCert, signedData, attestationSignature);
        verifyCertificateRequirements(attestationCert, credentialAttestation);
    }

    private List<Certificate> fetchCertsFromAttestationStmt(Map<String, Object> attStmt)
            throws WebAuthnException
    {
        List<Object> x509Certs = (List)attStmt.get("x5c");
        List<Certificate> x509decodeCerts = getDecodedX509Certs(x509Certs);
        return x509decodeCerts;
    }

    private Certificate getAttestationCertificate(List<Certificate> x509decodeCerts)
    {
        Certificate attestationCert = (Certificate)x509decodeCerts.get(0);
        return attestationCert;
    }

    public void verifyCertificateRequirements(Certificate attestationCert, CredentialAttestation credentialAttestation)
            throws WebAuthnException
    {
        X509Certificate x509Cert = (X509Certificate)attestationCert;

        LOG.debug("Attestation certificate: {}", x509Cert.toString());

        verifyCertificateVersion(x509Cert);
        LOG.debug("Certificate version verification successfull");

        verifyCertificateSubjectFields(x509Cert);
        LOG.debug("Certificate subject verification successfull");

        verifyCertificateExtensions(x509Cert, credentialAttestation);
        LOG.debug("Certificate extension verification successfull");
    }

    private void verifyCertificateVersion(X509Certificate x509Cert)
            throws WebAuthnException
    {
        if (x509Cert.getVersion() != 3)
        {
            String errorMsg = "Invalid certificate version. Expected 3 found [" + x509Cert.getVersion() + "]";
            logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
        }
    }

    private void verifyCertificateSubjectFields(X509Certificate x509Cert)
            throws WebAuthnException
    {
        try
        {
            X500Name x500name = new JcaX509CertificateHolder(x509Cert).getSubject();
            if (x500name.getRDNs(BCStyle.C).length == 0)
            {
                String errorMsg = "Missing Subject-C field";
                logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
            }
            else
            {
                RDN c = x500name.getRDNs(BCStyle.C)[0];
                String value = IETFUtils.valueToString(c.getFirst().getValue());
                if (!LocationConstants.COUNTRY_CODE_NAME_MAP.containsKey(value))
                {
                    String errorMsg = "Invalid field Subject-C value";
                    logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
                }
            }
            if (x500name.getRDNs(BCStyle.O).length == 0)
            {
                String errorMsg = "Missing Subject-O field";
                logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
            }
            if (x500name.getRDNs(BCStyle.OU).length == 0)
            {
                String errorMsg = "Missing Subject-OU field";
                logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
            }
            else
            {
                RDN ou = x500name.getRDNs(BCStyle.OU)[0];
                String value = IETFUtils.valueToString(ou.getFirst().getValue());
                if (!"Authenticator Attestation".equalsIgnoreCase(value))
                {
                    String errorMsg = "Invalid field Subject-OU value";
                    logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
                }
            }
            if (x500name.getRDNs(BCStyle.CN).length == 0)
            {
                String errorMsg = "Missing Subject-CN field";
                logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
            }
        }
        catch (Exception ex)
        {
            String errorMsg = ex.getMessage();
            logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED, ex);
        }
    }

    private void verifyCertificateExtensions(X509Certificate x509Cert, CredentialAttestation credentialAttestation)
            throws WebAuthnException
    {
        if (x509Cert.getBasicConstraints() != -1)
        {
            String errorMsg = "Basic Constraints extension MUST have the CA component set to false";
            logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
        }
        verifyAAGUIDInCertExtentsion(x509Cert, credentialAttestation);
    }
}

